"""
Gradio Web界面应用
提供友好的用户界面来使用AI动画生成系统
"""

import gradio as gr
import torch
import numpy as np
import logging
import os
from pathlib import Path
from typing import Optional, Tuple, List
import tempfile
import time

# 添加项目根目录到Python路径
import sys
project_root = Path(__file__).parent.parent
sys.path.insert(0, str(project_root))

from models.text2video import TextToVideoGenerator
from models.style_transfer import StyleTransferModel
from models.video_editing import VideoEditor
from utils.video_utils import VideoProcessor
from configs.config_manager import ConfigManager

# 设置日志
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)

class GradioApp:
    """Gradio应用类"""
    
    def __init__(self):
        """初始化应用"""
        self.config_manager = ConfigManager()
        self.config = self.config_manager.get_config()
        
        # 检查设备
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        logger.info(f"使用设备: {self.device}")
        
        # 初始化模型
        self._init_models()
        
        # 创建输出目录
        self.output_dir = Path("outputs/videos")
        self.output_dir.mkdir(parents=True, exist_ok=True)
        
        # 创建界面
        self.interface = self._create_interface()
    
    def _init_models(self):
        """初始化模型"""
        try:
            logger.info("正在初始化模型...")
            
            self.text2video_generator = TextToVideoGenerator(
                config=self.config['text2video'],
                device=self.device
            )
            
            self.style_transfer_model = StyleTransferModel(
                config=self.config['style_transfer'],
                device=self.device
            )
            
            self.video_editor = VideoEditor(
                config=self.config['video_editing'],
                device=self.device
            )
            
            self.video_processor = VideoProcessor()
            
            logger.info("模型初始化完成!")
            
        except Exception as e:
            logger.error(f"模型初始化失败: {e}")
            raise
    
    def _create_interface(self) -> gr.Interface:
        """创建Gradio界面"""
        
        # 定义样式
        css = """
        .gradio-container {
            max-width: 1200px;
            margin: 0 auto;
        }
        .header {
            text-align: center;
            margin-bottom: 2rem;
        }
        .header h1 {
            color: #2c3e50;
            font-size: 2.5rem;
            margin-bottom: 0.5rem;
        }
        .header p {
            color: #7f8c8d;
            font-size: 1.1rem;
        }
        """
        
        with gr.Blocks(css=css, title="AI动画生成系统") as interface:
            
            # 标题
            with gr.Row():
                with gr.Column():
                    gr.HTML("""
                    <div class="header">
                        <h1>🎬 AI动画生成系统</h1>
                        <p>基于大模型的智能动画生成工具</p>
                    </div>
                    """)
            
            # 主要功能标签页
            with gr.Tabs():
                
                # 文本到视频生成
                with gr.TabItem("🎨 文本生成动画"):
                    self._create_text2video_tab()
                
                # 风格迁移
                with gr.TabItem("🎭 风格迁移"):
                    self._create_style_transfer_tab()
                
                # 视频编辑
                with gr.TabItem("✂️ 视频编辑"):
                    self._create_video_edit_tab()
                
                # 批量处理
                with gr.TabItem("📦 批量处理"):
                    self._create_batch_tab()
            
            # 底部信息
            with gr.Row():
                with gr.Column():
                    gr.HTML("""
                    <div style="text-align: center; margin-top: 2rem; color: #7f8c8d;">
                        <p>AI动画生成系统 v1.0.0 | 基于大模型的智能创作工具</p>
                    </div>
                    """)
        
        return interface
    
    def _create_text2video_tab(self):
        """创建文本到视频生成标签页"""
        
        with gr.Row():
            with gr.Column(scale=1):
                # 输入参数
                prompt_input = gr.Textbox(
                    label="🎯 文本描述",
                    placeholder="请输入您想要生成的动画描述，例如：一只可爱的小猫在花园里玩耍",
                    lines=3,
                    max_lines=5
                )
                
                style_dropdown = gr.Dropdown(
                    choices=["realistic", "cartoon", "anime", "watercolor", "oil_painting", "sketch"],
                    value="realistic",
                    label="🎨 风格选择",
                    info="选择动画的艺术风格"
                )
                
                duration_slider = gr.Slider(
                    minimum=1.0,
                    maximum=10.0,
                    value=3.0,
                    step=0.5,
                    label="⏱️ 视频时长 (秒)",
                    info="设置生成视频的时长"
                )
                
                resolution_dropdown = gr.Dropdown(
                    choices=["256x256", "512x512", "768x768"],
                    value="256x256",
                    label="📐 分辨率",
                    info="选择视频分辨率"
                )
                
                steps_slider = gr.Slider(
                    minimum=20,
                    maximum=100,
                    value=50,
                    step=5,
                    label="🔄 生成步数",
                    info="更多的步数通常产生更好的质量，但需要更长时间"
                )
                
                generate_btn = gr.Button(
                    "🚀 生成动画",
                    variant="primary",
                    size="lg"
                )
            
            with gr.Column(scale=1):
                # 输出结果
                output_video = gr.Video(
                    label="🎬 生成的动画",
                    format="mp4"
                )
                
                status_text = gr.Textbox(
                    label="📊 生成状态",
                    interactive=False,
                    lines=2
                )
                
                # 示例提示词
                with gr.Accordion("💡 示例提示词", open=False):
                    gr.HTML("""
                    <div style="padding: 1rem;">
                        <h4>动物动画：</h4>
                        <ul>
                            <li>一只可爱的小猫在花园里玩耍</li>
                            <li>小狗在草地上奔跑追逐蝴蝶</li>
                            <li>熊猫在竹林中悠闲地吃竹子</li>
                        </ul>
                        
                        <h4>人物动画：</h4>
                        <ul>
                            <li>小女孩在花海中跳舞</li>
                            <li>宇航员在太空中漂浮</li>
                            <li>舞者在舞台上表演芭蕾</li>
                        </ul>
                        
                        <h4>场景动画：</h4>
                        <ul>
                            <li>日落时分的海滩波浪</li>
                            <li>森林中的瀑布流水</li>
                            <li>城市夜景中的车流</li>
                        </ul>
                    </div>
                    """)
        
        # 绑定生成函数
        generate_btn.click(
            fn=self._generate_video,
            inputs=[prompt_input, style_dropdown, duration_slider, resolution_dropdown, steps_slider],
            outputs=[output_video, status_text]
        )
    
    def _create_style_transfer_tab(self):
        """创建风格迁移标签页"""
        
        with gr.Row():
            with gr.Column(scale=1):
                # 输入视频
                input_video = gr.Video(
                    label="📹 输入视频",
                    format="mp4"
                )
                
                style_dropdown = gr.Dropdown(
                    choices=["cartoon", "anime", "watercolor", "oil_painting", "sketch"],
                    value="cartoon",
                    label="🎨 目标风格",
                    info="选择要应用的风格"
                )
                
                strength_slider = gr.Slider(
                    minimum=0.1,
                    maximum=1.0,
                    value=0.7,
                    step=0.1,
                    label="💪 风格强度",
                    info="控制风格迁移的强度"
                )
                
                transfer_btn = gr.Button(
                    "🎭 应用风格",
                    variant="primary",
                    size="lg"
                )
            
            with gr.Column(scale=1):
                # 输出结果
                output_video = gr.Video(
                    label="🎨 风格化视频",
                    format="mp4"
                )
                
                status_text = gr.Textbox(
                    label="📊 处理状态",
                    interactive=False,
                    lines=2
                )
                
                # 风格预览
                with gr.Accordion("🎨 风格预览", open=False):
                    gr.HTML("""
                    <div style="padding: 1rem;">
                        <h4>可用风格：</h4>
                        <ul>
                            <li><strong>卡通风格</strong>：明亮的卡通动画效果</li>
                            <li><strong>动漫风格</strong>：日本动漫风格</li>
                            <li><strong>水彩画风格</strong>：柔和的水彩画效果</li>
                            <li><strong>油画风格</strong>：厚重的油画效果</li>
                            <li><strong>素描风格</strong>：黑白素描效果</li>
                        </ul>
                    </div>
                    """)
        
        # 绑定风格迁移函数
        transfer_btn.click(
            fn=self._apply_style_transfer,
            inputs=[input_video, style_dropdown, strength_slider],
            outputs=[output_video, status_text]
        )
    
    def _create_video_edit_tab(self):
        """创建视频编辑标签页"""
        
        with gr.Row():
            with gr.Column(scale=1):
                # 输入视频
                input_video = gr.Video(
                    label="📹 输入视频",
                    format="mp4"
                )
                
                edit_prompt = gr.Textbox(
                    label="✏️ 编辑指令",
                    placeholder="请输入编辑指令，例如：增强视频质量、去除噪声、调整速度等",
                    lines=2
                )
                
                operation_dropdown = gr.Dropdown(
                    choices=["enhance", "stabilize", "denoise", "upscale", "color_correct", "speed"],
                    value="enhance",
                    label="🔧 编辑操作",
                    info="选择要执行的编辑操作"
                )
                
                edit_btn = gr.Button(
                    "✂️ 编辑视频",
                    variant="primary",
                    size="lg"
                )
            
            with gr.Column(scale=1):
                # 输出结果
                output_video = gr.Video(
                    label="✂️ 编辑后视频",
                    format="mp4"
                )
                
                status_text = gr.Textbox(
                    label="📊 编辑状态",
                    interactive=False,
                    lines=2
                )
                
                # 编辑操作说明
                with gr.Accordion("📖 编辑操作说明", open=False):
                    gr.HTML("""
                    <div style="padding: 1rem;">
                        <h4>可用操作：</h4>
                        <ul>
                            <li><strong>质量增强</strong>：提升视频整体质量</li>
                            <li><strong>视频稳定</strong>：减少视频抖动</li>
                            <li><strong>降噪</strong>：去除视频噪声</li>
                            <li><strong>超分辨率</strong>：提升视频分辨率</li>
                            <li><strong>颜色校正</strong>：自动颜色校正</li>
                            <li><strong>速度调整</strong>：调整播放速度</li>
                        </ul>
                    </div>
                    """)
        
        # 绑定编辑函数
        edit_btn.click(
            fn=self._edit_video,
            inputs=[input_video, edit_prompt, operation_dropdown],
            outputs=[output_video, status_text]
        )
    
    def _create_batch_tab(self):
        """创建批量处理标签页"""
        
        with gr.Row():
            with gr.Column(scale=1):
                # 批量输入
                input_files = gr.File(
                    file_count="multiple",
                    label="📁 选择视频文件",
                    file_types=["video"]
                )
                
                batch_operation = gr.Dropdown(
                    choices=["style_transfer", "enhance", "denoise"],
                    value="style_transfer",
                    label="🔧 批量操作",
                    info="选择要执行的批量操作"
                )
                
                batch_style = gr.Dropdown(
                    choices=["cartoon", "anime", "watercolor"],
                    value="cartoon",
                    label="🎨 目标风格",
                    visible=True
                )
                
                batch_btn = gr.Button(
                    "📦 开始批量处理",
                    variant="primary",
                    size="lg"
                )
            
            with gr.Column(scale=1):
                # 处理结果
                output_files = gr.File(
                    label="📁 处理结果",
                    file_count="multiple"
                )
                
                progress_text = gr.Textbox(
                    label="📊 处理进度",
                    interactive=False,
                    lines=3
                )
        
        # 绑定批量处理函数
        batch_btn.click(
            fn=self._batch_process,
            inputs=[input_files, batch_operation, batch_style],
            outputs=[output_files, progress_text]
        )
    
    def _generate_video(self, prompt: str, style: str, duration: float, 
                       resolution: str, steps: int) -> Tuple[str, str]:
        """生成视频"""
        try:
            if not prompt.strip():
                return None, "❌ 请输入文本描述"
            
            # 更新状态
            status = f"🚀 开始生成动画: {prompt[:50]}..."
            
            # 解析分辨率
            width, height = map(int, resolution.split('x'))
            
            # 生成视频
            video = self.text2video_generator.generate(
                prompt=prompt,
                duration=duration,
                style=style,
                num_inference_steps=steps
            )
            
            # 保存视频
            timestamp = int(time.time())
            output_path = self.output_dir / f"generated_{timestamp}.mp4"
            
            self.video_processor.save_video(video, str(output_path))
            
            status = f"✅ 动画生成完成! 保存至: {output_path.name}"
            
            return str(output_path), status
            
        except Exception as e:
            logger.error(f"生成视频失败: {e}")
            return None, f"❌ 生成失败: {str(e)}"
    
    def _apply_style_transfer(self, video_path: str, style: str, strength: float) -> Tuple[str, str]:
        """应用风格迁移"""
        try:
            if not video_path:
                return None, "❌ 请选择输入视频"
            
            status = f"🎭 正在应用 {style} 风格..."
            
            # 加载视频
            video = self.video_processor.load_video(video_path)
            
            # 应用风格迁移
            styled_video = self.style_transfer_model.transfer_style(video, style)
            
            # 保存结果
            timestamp = int(time.time())
            output_path = self.output_dir / f"styled_{style}_{timestamp}.mp4"
            
            self.video_processor.save_video(styled_video, str(output_path))
            
            status = f"✅ 风格迁移完成! 保存至: {output_path.name}"
            
            return str(output_path), status
            
        except Exception as e:
            logger.error(f"风格迁移失败: {e}")
            return None, f"❌ 风格迁移失败: {str(e)}"
    
    def _edit_video(self, video_path: str, edit_prompt: str, operation: str) -> Tuple[str, str]:
        """编辑视频"""
        try:
            if not video_path:
                return None, "❌ 请选择输入视频"
            
            status = f"✂️ 正在执行 {operation} 操作..."
            
            # 加载视频
            video = self.video_processor.load_video(video_path)
            
            # 编辑视频
            edited_video = self.video_editor.edit(video, edit_prompt)
            
            # 保存结果
            timestamp = int(time.time())
            output_path = self.output_dir / f"edited_{operation}_{timestamp}.mp4"
            
            self.video_processor.save_video(edited_video, str(output_path))
            
            status = f"✅ 视频编辑完成! 保存至: {output_path.name}"
            
            return str(output_path), status
            
        except Exception as e:
            logger.error(f"视频编辑失败: {e}")
            return None, f"❌ 视频编辑失败: {str(e)}"
    
    def _batch_process(self, files: List[str], operation: str, style: str) -> Tuple[List[str], str]:
        """批量处理"""
        try:
            if not files:
                return [], "❌ 请选择要处理的文件"
            
            status = f"📦 开始批量处理 {len(files)} 个文件..."
            results = []
            
            for i, file_path in enumerate(files):
                try:
                    status = f"📦 处理进度: {i+1}/{len(files)} - {Path(file_path).name}"
                    
                    if operation == "style_transfer":
                        video = self.video_processor.load_video(file_path)
                        processed_video = self.style_transfer_model.transfer_style(video, style)
                    elif operation == "enhance":
                        video = self.video_processor.load_video(file_path)
                        processed_video = self.video_editor._enhance_quality(video, {}, {})
                    elif operation == "denoise":
                        video = self.video_processor.load_video(file_path)
                        processed_video = self.video_editor._denoise_video(video, {}, {})
                    
                    # 保存结果
                    timestamp = int(time.time())
                    output_path = self.output_dir / f"batch_{operation}_{i}_{timestamp}.mp4"
                    self.video_processor.save_video(processed_video, str(output_path))
                    results.append(str(output_path))
                    
                except Exception as e:
                    logger.error(f"处理文件 {file_path} 失败: {e}")
                    continue
            
            status = f"✅ 批量处理完成! 成功处理 {len(results)} 个文件"
            
            return results, status
            
        except Exception as e:
            logger.error(f"批量处理失败: {e}")
            return [], f"❌ 批量处理失败: {str(e)}"
    
    def launch(self, **kwargs):
        """启动应用"""
        return self.interface.launch(**kwargs)

def main():
    """主函数"""
    try:
        app = GradioApp()
        app.launch(
            server_name="0.0.0.0",
            server_port=7860,
            share=True,
            debug=True
        )
    except Exception as e:
        logger.error(f"启动应用失败: {e}")
        raise

if __name__ == "__main__":
    main() 